package com.arvin.faceDetect.utils

import android.graphics.Bitmap
import android.graphics.Canvas
import android.graphics.Color
import android.graphics.Matrix
import android.util.Log
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.coroutineScope
import kotlinx.coroutines.withContext
import kotlin.math.max
import kotlin.math.min

/**
 * 图像预处理器
 * 用于对人脸图像进行预处理，包括：
 * 1. 光照归一化
 * 2. 对比度增强
 * 3. 直方图均衡化
 * 4. 图像尺寸调整
 */
object ImagePreprocessor {
    private const val TAG = "ImagePreprocessor"

    // 目标尺寸
    private const val TARGET_SIZE = 224

    /**
     * 处理人脸图像
     * @param bitmap 原始人脸图像
     * @param shouldNormalizeLighting 是否进行光照归一化
     * @param shouldEnhanceContrast 是否增强对比度
     * @param shouldEqualizeHistogram 是否进行直方图均衡化
     * @return 处理后的图像
     */
    suspend fun process(
        bitmap: Bitmap,
        shouldNormalizeLighting: Boolean = true,
        shouldEnhanceContrast: Boolean = true,
        shouldEqualizeHistogram: Boolean = true
    ): Bitmap = coroutineScope {
        try {
            // 1. 分析图像质量
            val qualityOptions = ImageQualityAnalyzer.analyzeImageQuality(bitmap)

            // 2. 调整图像大小
            val resizedBitmap = resizeBitmap(bitmap)

            // 3. 并行执行预处理步骤
            val processedBitmap = withContext(Dispatchers.Default) {
                var result = resizedBitmap

                // 根据质量分析结果决定是否执行各个预处理步骤
                if (qualityOptions.needLightingNormalization && shouldNormalizeLighting) {
                    result = normalizeLighting(result)
                }

                if (qualityOptions.needContrastEnhancement && shouldEnhanceContrast) {
                    result = enhanceContrast(result)
                }

                if (qualityOptions.needHistogramEqualization && shouldEqualizeHistogram) {
                    result = equalizeHistogram(result)
                }

                result
            }

            processedBitmap
        } catch (e: Exception) {
            Log.e(TAG, "图像预处理失败: ${e.message}")
            bitmap
        }
    }

    /**
     * 调整图像大小
     */
    private fun resizeBitmap(bitmap: Bitmap): Bitmap {
        val width = bitmap.width
        val height = bitmap.height

        if (width == TARGET_SIZE && height == TARGET_SIZE) {
            return bitmap
        }

        val scale = min(
            TARGET_SIZE.toFloat() / width,
            TARGET_SIZE.toFloat() / height
        )

        val matrix = Matrix().apply {
            postScale(scale, scale)
        }

        return Bitmap.createBitmap(
            bitmap,
            0,
            0,
            width,
            height,
            matrix,
            true
        )
    }

    /**
     * 光照归一化
     * 使用Gamma校正和亮度均衡
     */
    private fun normalizeLighting(bitmap: Bitmap): Bitmap {
        val width = bitmap.width
        val height = bitmap.height
        val pixels = IntArray(width * height)
        bitmap.getPixels(pixels, 0, width, 0, 0, width, height)

        // 计算平均亮度
        var totalBrightness = 0f
        for (pixel in pixels) {
            val r = Color.red(pixel)
            val g = Color.green(pixel)
            val b = Color.blue(pixel)
            totalBrightness += (r + g + b) / (3f * 255f)
        }
        val avgBrightness = totalBrightness / (width * height)

        // 目标亮度
        val targetBrightness = 0.5f

        // 计算亮度调整因子
        val factor = if (avgBrightness > 0) targetBrightness / avgBrightness else 1f

        // 创建新的Bitmap
        val resultBitmap = BitmapPool.getBitmap(width, height)
        val canvas = Canvas(resultBitmap)

        // 应用亮度调整
        for (i in pixels.indices) {
            val pixel = pixels[i]
            val r = (Color.red(pixel) * factor).toInt().coerceIn(0, 255)
            val g = (Color.green(pixel) * factor).toInt().coerceIn(0, 255)
            val b = (Color.blue(pixel) * factor).toInt().coerceIn(0, 255)
            pixels[i] = Color.rgb(r, g, b)
        }

        resultBitmap.setPixels(pixels, 0, width, 0, 0, width, height)
        return resultBitmap
    }

    /**
     * 对比度增强
     * 使用自适应对比度增强
     */
    private fun enhanceContrast(bitmap: Bitmap): Bitmap {
        val width = bitmap.width
        val height = bitmap.height
        val pixels = IntArray(width * height)
        bitmap.getPixels(pixels, 0, width, 0, 0, width, height)

        // 计算亮度范围
        var minBrightness = 1f
        var maxBrightness = 0f
        for (pixel in pixels) {
            val r = Color.red(pixel)
            val g = Color.green(pixel)
            val b = Color.blue(pixel)
            val brightness = (r + g + b) / (3f * 255f)
            minBrightness = min(minBrightness, brightness)
            maxBrightness = max(maxBrightness, brightness)
        }

        // 计算对比度调整因子
        val range = maxBrightness - minBrightness
        val factor = if (range > 0) 1f / range else 1f

        // 创建新的Bitmap
        val resultBitmap = BitmapPool.getBitmap(width, height)

        // 应用对比度增强
        for (i in pixels.indices) {
            val pixel = pixels[i]
            val r = Color.red(pixel)
            val g = Color.green(pixel)
            val b = Color.blue(pixel)

            val newR = ((r / 255f - minBrightness) * factor * 255f).toInt().coerceIn(0, 255)
            val newG = ((g / 255f - minBrightness) * factor * 255f).toInt().coerceIn(0, 255)
            val newB = ((b / 255f - minBrightness) * factor * 255f).toInt().coerceIn(0, 255)

            pixels[i] = Color.rgb(newR, newG, newB)
        }

        resultBitmap.setPixels(pixels, 0, width, 0, 0, width, height)
        return resultBitmap
    }

    /**
     * 直方图均衡化
     * 使用RGB通道分别进行均衡化
     */
    private fun equalizeHistogram(bitmap: Bitmap): Bitmap {
        val width = bitmap.width
        val height = bitmap.height
        val pixels = IntArray(width * height)
        bitmap.getPixels(pixels, 0, width, 0, 0, width, height)

        // 计算直方图
        val histogram = IntArray(256)
        for (pixel in pixels) {
            val r = Color.red(pixel)
            val g = Color.green(pixel)
            val b = Color.blue(pixel)
            val gray = ((r + g + b) / 3).toInt()
            histogram[gray]++
        }

        // 计算累积直方图
        val cdf = IntArray(256)
        var sum = 0
        for (i in 0..255) {
            sum += histogram[i]
            cdf[i] = sum
        }

        // 归一化累积直方图
        val total = width * height
        val normalizedCdf = FloatArray(256)
        for (i in 0..255) {
            normalizedCdf[i] = cdf[i].toFloat() / total
        }

        // 创建新的Bitmap
        val resultBitmap = BitmapPool.getBitmap(width, height)

        // 应用直方图均衡化
        for (i in pixels.indices) {
            val pixel = pixels[i]
            val r = Color.red(pixel)
            val g = Color.green(pixel)
            val b = Color.blue(pixel)

            val gray = ((r + g + b) / 3).toInt()
            val newGray = (normalizedCdf[gray] * 255).toInt().coerceIn(0, 255)

            pixels[i] = Color.rgb(newGray, newGray, newGray)
        }

        resultBitmap.setPixels(pixels, 0, width, 0, 0, width, height)
        return resultBitmap
    }
} 